import requests
import ffmpeg
import os
import re
import shutil


output_texts_file = r'D:\Documents\work\video\gitee\3\chdoc.txt' # 纯文本中文字幕
output_subtitles_file = r'D:\Documents\work\video\gitee\3\ch.txt' # 带时间戳中文字幕

wave_FOLDER = 'output_wav_files'
wave_time = 'ch_d_time.txt' # 读取中文音频时长
new_ch = 'updated_ch.txt' # 要读取的 ch.txt 文件


# 定义下载音频的函数
def getAudio(savename, url):
    response = requests.get(url)
    if response.status_code == 200:
        with open(savename, "wb") as file:
            file.write(response.content)
        print(f"音频文件 {savename} 已成功下载到本地。")
        return True
    else:
        print(f"下载失败，状态码: {response.status_code}")
        return False

# 使用 ChatTTS 生成音频
def chattts(msg, audioname):
    res = requests.post('http://192.168.88.251:9966/tts', data={
        "text": msg,
        "prompt": "",
        "voice": "7869",
        "temperature": 0.3,
        "top_p": 0.7,
        "top_k": 20,
        "skip_refine": 0,
        "custom_voice": 0,
        "is_split": 1
    })
    if res.status_code == 200:
        url = res.json().get('url')
        if url:
            getAudio(audioname, url)
        else:
            print("未能获取音频文件的 URL。")
    else:
        print(f"请求失败，状态码: {res.status_code}")

# 生成字幕音频文件
def process_subtitles(output_texts_file, OUTPUT_FOLDER):
    with open(output_texts_file, 'r', encoding='utf-8') as file:
        lines = file.readlines()
    for index, line in enumerate(lines):
        line = line.strip()
        if line:
            audioname = f'{OUTPUT_FOLDER}/line_{index + 1}.wav'
            chattts(line, audioname)

# 获取音频时长并保存
def get_audio_durations(OUTPUT_FOLDER, audio_durations_file):
    audio_durations = []
    for filename in os.listdir(OUTPUT_FOLDER):
        if filename.endswith('.wav') and filename.startswith('line_'):
            match = re.search(r'line_(\d+)', filename)
            if match:
                index = int(match.group(1))
                file_path = os.path.join(OUTPUT_FOLDER, filename)
                try:
                    probe = ffmpeg.probe(file_path)
                    duration_sec = float(probe['format']['duration'])
                    duration_ms = int(duration_sec * 1000)
                    audio_durations.append((index, filename, duration_ms))
                except ffmpeg.Error as e:
                    print(f"无法处理文件 {filename}: {e}")
    audio_durations.sort()
    with open(audio_durations_file, 'w') as f:
        for _, filename, duration_ms in audio_durations:
            f.write(f"{filename}: {duration_ms} ms\n")
    print(f"音频时长已写入 {audio_durations_file}")

# 更新字幕文件
def update_ch_file(output_subtitles_file, audio_durations_file, output_file):
    durations = {}
    with open(audio_durations_file, 'r', encoding='utf-8') as f:
        for line in f:
            match = re.match(r'line_(\d+).wav: (\d+) ms', line.strip())
            if match:
                index = int(match.group(1))
                duration = int(match.group(2))
                durations[index] = duration
    updated_lines = []
    with open(output_subtitles_file, 'r', encoding='utf-8') as f:
        for line in f:
            match = re.match(r'(\d+): \[(.*)\]', line.strip())
            if match:
                index = int(match.group(1))
                if index in durations:
                    updated_content = f"{match.group(2)}, {durations[index]}"
                    updated_line = f"{index}: [{updated_content}]\n"
                    updated_lines.append(updated_line)
    with open(output_file, 'w', encoding='utf-8') as f:
        f.writelines(updated_lines)
    print(f"更新后的内容已写入 {output_file}")

# 如果文件夹已存在，则删除
if os.path.exists(wave_FOLDER):
    shutil.rmtree(wave_FOLDER)  # 删除整个文件夹及其内容

# 重新创建空文件夹
os.makedirs(wave_FOLDER)

process_subtitles(output_texts_file, wave_FOLDER)  # 处理字幕文件并生成音频
get_audio_durations(wave_FOLDER, wave_time)  # 获取音频时长并保存
update_ch_file(output_subtitles_file, wave_time, new_ch) # 更新字幕文件
  